////////////////////////////////////////////////////////////////////////////////
// Includes
////////////////////////////////////////////////////////////////////////////////

#include "StructureEngine.hpp"

#include <wx/string.h>
#include <wx/colour.h>
#include <wx/gdicmn.h>
#include <cmath>
#include "Exception.hpp"

using namespace std;
using namespace Types;
using namespace Exceptions;

////////////////////////////////////////////////////////////////////////////////
// Namespaces
////////////////////////////////////////////////////////////////////////////////

namespace SC {

////////////////////////////////////////////////////////////////////////////////
// Implementation
////////////////////////////////////////////////////////////////////////////////

////////////////////////////////////////////////////////////////////////////////
// Public functions
////////////////////////////////////////////////////////////////////////////////

void StructureEngine::Initialise(const bool init_all) {
	if (this->graph == NULL) return;

	// Forces applied
	this->forces.resize(this->graph->GetNbNodes(), Vector3D<float>(0.f, 0.f, 0.f));

	// Clear interactions
	this->interaction_links.clear();

	// Eventually clear the selected systems
	if (init_all) {
		this->selected_systems.clear();
		this->nodes_selected_state.clear();
	}

	// Node init
	for (unsigned int n=0; n<this->graph->GetNbNodes(); ++n) {
		// Get node reference
		Node<NData,AData> & node = this->graph->Node(n);
		// Sets up initial node velocities to 0
		node.Data().AddData().struct_velocity = Vector3D<float>(0.f, 0.f, 0.f);
		// Sets up arcs
		for (unsigned int a=0; a<node.GetNbArcs(); ++a) {
			node.Arc(a).Data().AddData().struct_optimum = 1.f;
			node.Arc(a).Data().AddData().struct_strength = 1.f;
		}
	}

	// Node selected state
	this->nodes_selected_state.resize(this->graph->GetNbNodes());
	this->SetSelectedSystems(this->selected_systems);

	// Set skeleton's atoms' mass
	this->DistributeSkeletonMass(1.f);

	// Initialise locations
	this->ResetLocations();
}

void StructureEngine::Clear() {
	this->graph = NULL;
	this->forces.clear();
	this->interaction_links.clear();
	this->selected_systems.clear();
	this->nodes_selected_state.clear();
}

void StructureEngine::SetAnchor(const unsigned int value) {
	if ((this->graph != NULL) && (value < this->graph->GetNbNodes()))
		this->anchor = value;
}

void StructureEngine::ResetLocations() {
	if (this->graph == NULL) return;
	// Initialise all locations at 0
	for (unsigned int n=0; n<this->graph->GetNbNodes(); ++n) {
		this->graph->Node(n).Data().AddData().struct_location = Point3D<float>(0.f, 0.f, 0.f);
		this->graph->Node(n).Data().AddData().struct_velocity = Vector3D<float>(0.f, 0.f, 0.f);
	}
}

void StructureEngine::SetFDAlg(const FDA value) {
	this->fd_alg = value;
	for (unsigned int i=0; i<this->forces.size(); ++i)
		this->forces[i] = Vector3D<float>(0.f, 0.f, 0.f);
}

void StructureEngine::Enable3D(const bool value) {
	if (this->enable_3D == value) return;
	this->enable_3D = value;
	if (this->graph == NULL) return;
	if (value) {
		for (unsigned int n=0; n<this->graph->GetNbNodes(); ++n)
			this->graph->Node(n).Data().AddData().struct_location.z = 0.02f*frand()-0.01f;
	}
	else {
		for (unsigned int n=0; n<this->graph->GetNbNodes(); ++n)
			this->graph->Node(n).Data().AddData().struct_location.z = 0.f;
	}
}

void StructureEngine::SetUsageOnZScale(const float value) {
	this->usage_on_z_scale = value;
	if ((this->graph != NULL) && (!this->enable_3D) && (value <= 0.f)) {
		for (unsigned int n=0; n<this->graph->GetNbNodes(); ++n)
			this->graph->Node(n).Data().AddData().struct_location.z = 0.f;
	}
}

void StructureEngine::SetSelectedSystems(std::vector<int> & value) {
	if ((this->graph == NULL) || (this->nodes_selected_state.size() != this->graph->GetNbNodes())) return;
	this->selected_systems = value;
	// Set nodes selected state
	for (unsigned int n=0; n<this->graph->GetNbNodes(); ++n) {
		if (!this->selected_systems.empty() &&
			(find(this->selected_systems.begin(), this->selected_systems.end(), this->graph->GetNode(n).GetData().GetColourIndex()) == this->selected_systems.end()))
			this->nodes_selected_state[n] = false;
		else
			this->nodes_selected_state[n] = true;
	}
}

void StructureEngine::DistributeSkeletonMass(const float factor) {
	if (factor == 1.f) {
		for (unsigned int n=0; n<this->graph->GetNbNodes(); ++n)
			this->graph->Node(n).Data().AddData().struct_mass = 1.f;
	}
	else {
		// Initialise mass with respect to hierarchy
		vector<unsigned int> nodes(1, this->root_idx);
		vector<float> values(1, 1.f);
		unsigned int next_index = 0;
		while (next_index < nodes.size()) {
			// Set current node
			Node<NData,AData> & node = this->graph->Node(nodes[next_index]);
			node.Data().AddData().struct_mass = values[next_index];
			// Add more nodes
			for (unsigned int a=0; a<node.GetNbArcs(); ++a) {
				vector<unsigned int>::iterator itor = find(nodes.begin(), nodes.end(), node.GetArc(a).GetTargetNodeIndex());
				if (itor == nodes.end()) {
					nodes.push_back(node.GetArc(a).GetTargetNodeIndex());
					values.push_back(values[next_index] * factor);
				}
				else {
					unsigned int idx = itor-nodes.begin();
					if (values[idx] < values[next_index] * factor) {
						values[idx] = values[next_index] * factor;
						// Ammend node if needed
						if (idx < next_index)
							this->graph->Node(nodes[idx]).Data().AddData().struct_mass = values[idx];
					}
				}
			}
			// Next node
			++next_index;
		}
	}
}

void StructureEngine::RefreshNode(const unsigned int idx) {
	if (!this->selected_systems.empty()) {
		if (find(this->selected_systems.begin(), this->selected_systems.end(), this->graph->GetNode(idx).GetData().GetColourIndex()) == this->selected_systems.end())
			this->nodes_selected_state[idx] = false;
		else
			this->nodes_selected_state[idx] = true;
	}
}

const bool StructureEngine::Update(float & stability) {
	if (this->graph == NULL) return false;

	// Set Z-value
	if ((!this->enable_3D) && (this->usage_on_z_scale > 0.f)) {
		for (unsigned int n=0; n<this->graph->GetNbNodes(); ++n)
			this->graph->Node(n).Data().AddData().struct_location.z = this->graph->Node(n).Data().AddData().usage * this->usage_on_z_scale;
	}

	// Apply a force-directed algorithm
	switch (this->fd_alg) {
		case NormalFDA: {
			this->fda();
			break;
		}
		case HierarchicalFDA: {
			this->hfda();
			break;
		}
		default: throw Exception("Error: No structure layout algorithm selected");
	}

	// Compute spring tension due to interaction links
	for (unsigned int i=0; i<this->interaction_links.size(); ++i) {
		// Get nodes indices
		unsigned int n1 = this->interaction_links[i].GetFirst(), n2 = this->interaction_links[i].GetSecond();
		// If nodes are not to be considered, skip
		if (!this->nodes_selected_state[n1] || !this->nodes_selected_state[n2]) continue;
		// Update forces
		Vector3D<float> hf = 	this->hooke_attraction(
									this->graph->GetNode(n1).GetData().GetAddData().struct_location,
									this->graph->GetNode(n2).GetData().GetAddData().struct_location,
									this->interaction_links[i].GetThird(), 0.5f
								);
		if (this->anchor == n1)
			this->forces[n2] -= 2.f * hf;
		else if (this->anchor == n2)
			this->forces[n1] += 2.f * hf;
		else {
			this->forces[n1] += hf;
			this->forces[n2] -= hf;
		}
	}

	// Running sum of total kinetic energy over all particles
	float total_kinetic_energy = 0.f;

	// Apply the forces to the nodes
	for (unsigned int n=0; n<this->graph->GetNbNodes(); ++n) {
		// If node is not to be considered, skip
		if (!this->nodes_selected_state[n]) continue;
		// Get node reference
		Node<NData,AData> & node = this->graph->Node(n);
		// Update node velocity
		node.Data().AddData().struct_velocity += (this->timestep/node.Data().AddData().struct_mass) * forces[n];
		node.Data().AddData().struct_velocity *= this->damping;
		if (node.Data().AddData().struct_velocity.Norm() > 10.f) {
			node.Data().AddData().struct_velocity.Normalise();
			node.Data().AddData().struct_velocity *= 10.f;
		}
		// Reset force for this node
		forces[n] = Vector3D<float>(0.f, 0.f, 0.f);
		// Update total kinetic energy
		total_kinetic_energy = total_kinetic_energy + node.Data().AddData().struct_mass * node.Data().AddData().struct_velocity.SquaredNorm();
	}

	// If the energy is above the threshold the simulation is still running
	if (total_kinetic_energy > energy_threshold) {

		// Get anchor system motion vector
		//Vector3D<float> shift = this->graph->Node(this->anchor).Data().AddData().struct_velocity;

		// Update nodes locations
		for (unsigned int n=0; n<this->graph->GetNbNodes(); ++n) {
			// If node is not to be considered, skip
			if (!this->nodes_selected_state[n]) continue;
			// Get node reference
			Node<NData,AData> & node = this->graph->Node(n);
			// Update location
			node.Data().AddData().struct_location = node.Data().AddData().struct_location + this->timestep * (node.Data().AddData().struct_velocity);// - shift);
			if ((!this->enable_3D) && (this->usage_on_z_scale > 0.f) && (node.Data().AddData().struct_location.z > this->usage_on_z_scale))
				node.Data().AddData().struct_location.z = this->usage_on_z_scale;
			// Reset speed
			node.Data().AddData().struct_velocity = Vector3D<float>(0.f, 0.f, 0.f);
		}

		// Set the stability measure within [0,1]
		stability = total_kinetic_energy;
		stability = 1.f - MIN(stability,1.f);

		// Returns true to indicate that motion is still in progress
		return true;
	}
	// Otherwise it has stopped moving
	else return false;

}

////////////////////////////////////////////////////////////////////////////////
// Protected functions
////////////////////////////////////////////////////////////////////////////////

void StructureEngine::fda() {

	vector< pair<unsigned int, unsigned int> > idx_pairs;

	// Compute the total forces for all nodes
	for (unsigned int n=0; n<this->graph->GetNbNodes(); ++n) {

		// If node is not to be considered, skip
		if (!this->nodes_selected_state[n]) continue;

		// Get node reference
		Node<NData,AData> & node = this->graph->Node(n);

		// 1: Compute repulsion between each pair of nodes
		for (unsigned int n2=n+1; n2<this->graph->GetNbNodes(); ++n2) {
			// If node is not to be considered, skip
			if (!this->nodes_selected_state[n2]) continue;
			Vector3D<float> cf = this->coulomb_repulsion(node, this->graph->GetNode(n2));
			if (this->anchor == n)
				this->forces[n2] -= 2.f * cf;
			else if (this->anchor == n2)
				this->forces[n] += 2.f * cf;
			else {
				this->forces[n] += cf;
				this->forces[n2] -= cf;
			}
		}

		// 2: Compute spring tension for scopes
		vector<unsigned int> idx_stack;
		this->fda_att(n, n, idx_stack, idx_pairs);

	}
}

void StructureEngine::fda_att(
	unsigned int n,
	unsigned int nc,
	vector<unsigned int> & idx_stack,
	vector< pair<unsigned int, unsigned int> > & idx_pairs
) {
	// Check for recursive hierarchy
	if (find(idx_stack.begin(), idx_stack.end(), nc) != idx_stack.end()) return;
	// Check all subnodes of current node
	Node<NData,AData> & cur_node = this->graph->Node(nc);
	for (unsigned int a=0; a<cur_node.GetNbArcs(); ++a) {
		unsigned int n2 = cur_node.GetArc(a).GetTargetNodeIndex();
		// Create pair and check it was not already processed, if so jump
		pair<unsigned int, unsigned int> idx_pair(n,n2);
		if (find(idx_pairs.begin(), idx_pairs.end(), idx_pair) != idx_pairs.end()) continue;
		// If node is to be considered, update forces between the nodes
		if (this->nodes_selected_state[n2]) {
			Node<NData,AData> & root_node = this->graph->Node(n);
			// Update forces
			Vector3D<float> hf =	this->hooke_attraction(
										root_node.GetData().GetAddData().struct_location,
										this->graph->GetNode(n2).GetData().GetAddData().struct_location,
										root_node.GetArc(a).GetData().GetAddData().struct_optimum,
										root_node.GetArc(a).GetData().GetAddData().struct_strength
									);
			if (this->anchor == n)
				this->forces[n2] -= 2.f*hf;
			else if (this->anchor == n2)
				this->forces[n] += 2.f*hf;
			else {
				this->forces[n] += hf;
				this->forces[n2] -= hf;
			}
			// Add pair
			idx_pairs.push_back(idx_pair);
		}
		// Otherwise check its subnodes
		else {
			idx_stack.push_back(nc);
			this->fda_att(n, n2, idx_stack, idx_pairs);
			idx_stack.pop_back();
		}
	}
}


void StructureEngine::hfda() {
	// Apply hierarchical fda
	vector<unsigned int> path;
	this->hfda_step(this->root_idx, path);
}

void StructureEngine::hfda_step(const unsigned int node_idx, vector<unsigned int> & path) {

	// Check that node was not processed already
	if (find(path.begin(), path.end(), node_idx) != path.end()) return;

	// Add node as proccessed
	path.push_back(node_idx);

	// Get node reference
	const Node<NData,AData> & node = this->graph->GetNode(node_idx);

	// If node is to be considered, skip
	if (this->nodes_selected_state[node_idx]) {

		// For each node's sub node
		for (unsigned int a=0; a<node.GetNbArcs(); ++a) {

			// Get sub node index
			unsigned int sub_idx = node.GetArc(a).GetTargetNodeIndex();
			// If node is not to be considered, skip
			if (!this->nodes_selected_state[sub_idx]) continue;
			// Get sub node reference
			const Node<NData,AData> & sub_node = this->graph->GetNode(sub_idx);
			// Compute repulsion of the path nodes onto the sub node
			for (unsigned int p=0; p<path.size(); ++p)
				this->forces[sub_idx] += this->coulomb_repulsion(sub_node, this->graph->GetNode(path[p]));
			// Compute spring tension of the node onto the sub node
			this->forces[sub_idx] += 	this->hooke_attraction(
											sub_node.GetData().GetAddData().struct_location,
											node.GetData().GetAddData().struct_location,
											node.GetArc(a).GetData().GetAddData().struct_optimum,
											node.GetArc(a).GetData().GetAddData().struct_strength
										);
			
			// For each fellow node (other sub node of root)
			for (unsigned int a2=a+1; a2<node.GetNbArcs(); ++a2) {
				// Get fellow node index
				unsigned int fel_idx = node.GetArc(a2).GetTargetNodeIndex();
				// If node is not to be considered, skip
				if (!this->nodes_selected_state[fel_idx]) continue;
				// Get fellow node reference
				Node<NData,AData> & fel_node = this->graph->Node(fel_idx);
				// Compute repulsion between node and fellow node
				Vector3D<float> cf = this->coulomb_repulsion(sub_node, fel_node);
				this->forces[sub_idx] += cf;
				this->forces[fel_idx] -= cf;

			}

			// Apply sub node's sub nodes the same force to move them along
			this->shift_sub_nodes_force(sub_idx, this->forces[sub_idx], path);

		}

	}

	// Process each sub node
	for (unsigned int a=0; a<node.GetNbArcs(); ++a) {
		// Get sub node index
		unsigned int sub_idx = node.GetArc(a).GetTargetNodeIndex();
		// Apply algorithm to sub node
		this->hfda_step(sub_idx, path);
	}

	// Remove node as proccessed
	path.pop_back();

}

void StructureEngine::shift_sub_nodes_force(const unsigned int node_idx, const Vector3D<float> & shift, vector<unsigned int> & path) {

	// Check that node was not processed already
	if (find(path.begin(), path.end(), node_idx) != path.end()) return;

	// Add node as proccessed
	path.push_back(node_idx);

	// Get node reference
	const Node<NData,AData> & node = this->graph->GetNode(node_idx);

	// For each node's sub node
	for (unsigned int a=0; a<node.GetNbArcs(); ++a) {
		// Get sub node index
		unsigned int sub_idx = node.GetArc(a).GetTargetNodeIndex();
		// Apply force if nodes are to be considered
		if (this->nodes_selected_state[node_idx] && this->nodes_selected_state[sub_idx])
			this->forces[sub_idx] += shift;
		// Apply it to sub nodes
		this->shift_sub_nodes_force(sub_idx, shift, path);
	}

	// Remove node as proccessed
	path.pop_back();
}

const Vector3D<float> StructureEngine::coulomb_repulsion(
	const Node<NData,AData> & n1,
	const Node<NData,AData> & n2
) {
	Vector3D<float> dist = n1.GetData().GetAddData().struct_location - n2.GetData().GetAddData().struct_location;
	while (dist.IsNull()) dist += Vector3D<float>(0.02f*frand()-0.01f, 0.02f*frand()-0.01f, this->enable_3D?(0.02f*frand()-0.01f):0.f);
	float f = this->repulsion_power * ( (n1.GetData().GetAddData().struct_mass * n2.GetData().GetAddData().struct_mass) / dist.Norm() );
	return dist.GetNormalised() * f;
}

const Vector3D<float> StructureEngine::hooke_attraction(
	const Point3D<float> & p1,
	const Point3D<float> & p2,
	const float optimum,
	const float strength
) {
	Vector3D<float> dist = p2 - p1;
	float f = (dist.Norm() - optimum) * this->attraction_power;
	while (dist.IsNull()) dist += Vector3D<float>(0.02f*frand()-0.01f, 0.02f*frand()-0.01f, this->enable_3D?(0.02f*frand()-0.01f):0.f);
	if (f >= 0.f) return dist.GetNormalised() * pow(f, strength);
	else return dist.GetNormalised() * -pow(-f, strength);
}

////////////////////////////////////////////////////////////////////////////////
// Closing Namespaces
////////////////////////////////////////////////////////////////////////////////

}

////////////////////////////////////////////////////////////////////////////////
